import json
import re

from llm import Embedding, enc
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


class SampleGenerator:
    def __init__(self, samples, func):
        self.samples = samples
        self.func = func
        self.embedding_matrix = []
        self.transformer_data()
        self.embedding_sample()

    def transformer_data(self):
        # 处理后应该有 embedding_content 和 return_content
        self.samples = [self.func(i) for i in self.samples]

    def embedding_sample(self):
        embedding_content_list = [i['embedding_content'] for i in self.samples]

        self.embedding_matrix = []

        batch_size = 64
        for i in range(0, len(embedding_content_list), batch_size):
            batch_content_list = embedding_content_list[i:i + batch_size]
            self.embedding_matrix.extend(Embedding(batch_content_list))

    def get_sample(self, content, top_n=3):
        embedding = Embedding(content)
        embedding = np.array(embedding).reshape(1, -1)
        similarities = cosine_similarity(embedding, self.embedding_matrix)
        top_indices = similarities.argsort()[0][-top_n:][::-1]
        return [self.samples[i]['return_content'] for i in top_indices[:top_n]]

    def add_sample(self, sample):
        sample = self.func(sample)
        self.samples.append(sample)
        self.embedding_matrix.append(Embedding(sample))


def func(x):
    return x


def get_code_sample():
    li = []
    with open('./sample/code9.json', encoding='utf8') as f:
        data = json.loads(f.read())
        for i in data:
            i['code'] = re.sub( r'```[^`]*$', '```', i['code'], flags=re.DOTALL)
            li.append({"embedding_content": i['question'],
               "return_content": [{"role":"user","content":i['question']},
                                  {"role":"assistant","content":i['code']}]})
    return SampleGenerator(li, func)



def get_deepsea_sample():
    # with open('./sample/deepsea6.json', encoding='utf8') as f:
    #     li = json.loads(f.read())
    li = []
    with open('./sample/deepsea7.jsonl', encoding='utf8') as f:
        for line in f:
            data = json.loads(line)
            li.append({"embedding_content": data['question'],
               "return_content": data['messages']})
    return SampleGenerator(li, func)


def get_navigate_sample():
    with open('./sample/navigate6.json', encoding='utf8') as f:
        li = json.loads(f.read())
    return SampleGenerator(li, func)

def get_flow_sample():
    with open('./sample/flow.jsonl', encoding='utf8') as f:
        li = []
        for line in f:
            data = json.loads(line)
            if len(data['messages']) < 11:
                li.append({'embedding_content': data['question'],
                           'return_content': data['messages']})

    return SampleGenerator(li, func)
